package queue_stack.leetcode.medium;

import java.util.*;

/**
 * @author bruin_du
 * @description 和最小的 k 个数对
 * @date 2022/8/16 23:23
 **/
class OfferII061_KSmallestPairs {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        //方法一：暴力加大根堆
        method1(nums1, nums2, k);

        //方法二：多路归并加小跟堆！！
        return method2(nums1, nums2, k);
    }

    private List<List<Integer>> method1(int[] nums1, int[] nums2, int k) {
        Queue<int[]> maxHeap = new PriorityQueue<>((p1, p2)
                -> p2[0] + p2[1] - p1[0] - p1[1]);
        for (int i = 0; i < Math.min(k, nums1.length); ++i) {
            for (int j = 0; j < Math.min(k, nums2.length); ++j) {
                if (maxHeap.size() >= k) {
                    int[] root = maxHeap.peek();
                    int sum = root[0] + root[1];
                    if (sum > nums1[i] + nums2[j]) {
                        maxHeap.poll();
                        maxHeap.offer(new int[] {nums1[i], nums2[j]});
                    } else {
                        break;
                    }
                } else if (maxHeap.size() < k){
                    maxHeap.offer(new int[] {nums1[i], nums2[j]});
                }
            }
        }
        List<List<Integer>> result = new LinkedList<>();
        while (!maxHeap.isEmpty()) {
            int[] val = maxHeap.poll();
            result.add(Arrays.asList(val[0], val[1]));
        }

        return result;
    }

    private List<List<Integer>> method2(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> lists = new ArrayList<>();
        Queue<int[]> queue = new PriorityQueue<>((o1, o2) -> {
            return nums1[o1[0]] + nums2[o1[1]] - nums1[o2[0]] - nums2[o2[1]];
        });
        for (int i = 0; i < Math.min(nums1.length, k); i++) {
            queue.offer(new int[]{i, 0});
        }
        while (k-- != 0 && !queue.isEmpty()) {
            //拿到小根堆的堆顶
            int[] arr = queue.poll();
            List<Integer> list = new ArrayList<>();
            list.add(nums1[arr[0]]);
            list.add(nums2[arr[1]]);
            lists.add(list);
            if (arr[1] + 1 < nums2.length) {
                //这里每次都会根据数组的值相加进行排序的，所以下一次拿到的坐标的两个数组的值相加就是最小的
                queue.offer(new int[]{arr[0], arr[1] + 1});
            }
        }
        return lists;
    }

    public static void main(String[] args) {
        int[] arr1 = new int[]{1, 7, 8};
        int[] arr2 = new int[]{2, 4, 6};
        OfferII061_KSmallestPairs offer = new OfferII061_KSmallestPairs();
        offer.kSmallestPairs(arr1, arr2, 5);
    }
}
